"""
Monkeypatches google.appengine.ext.db to use cached versions for entities
fetched by key (or key name, or id): db.get, db.put and db.delete are replaced
by cacheable versions.

Based on https://gist.github.com/43429
"""
from google.appengine.api import datastore
from google.appengine.api import namespace_manager
from google.appengine.ext import db

from tipfy.appengine.db import to_key


class DatastoreCache(object):
    _cache = {}

    def __init__(self):
        self.cached_get_count = 0
        self.uncached_get_count = 0

    def get_cache(self):
        """Returns a cache for the current namespace."""
        ns = namespace_manager.get_namespace()
        if ns not in self._cache:
            self._cache[ns] = {}

        return self._cache[ns]

    def get(self, keys, **kwargs):
        if db.is_in_transaction():
            # In transactions, don't get from cache at all.
            return db_get(keys, **kwargs)

        cache = self.get_cache()
        keys, multiple = datastore.NormalizeAndTypeCheckKeys(keys)

        if multiple:
            # Split into cached and uncached.
            cond_func = lambda x: str(x) in cache
            cached_keys, uncached_keys = split_by_cond(keys, cond_func)
            cached_models = [cache[str(x)] for x in cached_keys]

            self.cached_get_count += len(cached_keys)
            self.uncached_get_count += len(uncached_keys)

            if uncached_keys:
                # Fetch uncached.
                uncached_models = db_get(uncached_keys, **kwargs)

                # Update cache.
                cache.update((str(m.key()), m) for m in uncached_models if m)

                # Construct return list.
                res = list(join_by_cond(keys, cached_models, uncached_models,
                    cond_func))
            else:
                res = cached_models
        else:
            key = str(keys[0])
            if key not in cache:
                res = db_get(keys[0], **kwargs)
                if res:
                    cache[key] = res

                self.uncached_get_count += 1
            else:
                res = cache[key]
                self.cached_get_count += 1

        return res

    def put(self, models, **kwargs):
        cache = self.get_cache()
        keys = db_put(models, **kwargs)

        if db.is_in_transaction():
            # In transactions, delete from the cache, since we don't know if
            # it will be committed or rolled back.
            for key in keys:
                key = str(key)
                if key in cache:
                    del cache[key]
        else:
            cache.update((str(k), m) for k, m in zip(list(keys), list(models)))

        return keys

    def delete(self, models, **kwargs):
        cache = self.get_cache()

        if isinstance(models, (basestring, db.Model, db.Key)):
            models = [models]

        for key in to_key(models):
            key = str(key)
            if key in cache:
                del cache[key]

        return db_delete(models, **kwargs)


def split_by_cond(l, func):
    a = list()
    b = list()
    for item in l:
        if func(item):
            a.append(item)
        else:
            b.append(item)

    return (a,b)


def join_by_cond(l, a, b, func):
    ret = list()
    ia = iter(a)
    ib = iter(b)
    for item in l:
        if func(item):
            yield ia.next()
        else:
            yield ib.next()


cache = DatastoreCache()

# Backup and monkeypatch db.get, db.put and db.delete.
db_get = db.get
db_put = db.put
db_delete = db.delete

db.get = cache.get
db.put = cache.put
db.delete = cache.delete
